Add per-execution memory budgets to RVM - #792

Maksym (@maksym-mishchenko) Thanks for doing this very useful feature! Overall looks good to me. Copilot reviews found some interesting cases that are worth addressing.

Maksym (Maksym (@maksym-mishchenko)) Thanks for doing this very useful feature! Overall looks good to me. Copilot reviews found some interesting cases that are worth addressing.

Hi Anand Krishnamoorthi (@anakrish), thanks for the thorough review. I addressed comments, I kept the two API suggestions unchanged for the reasons explained in their threads. Could you please take another look when you have time?

Maksym (Maksym (Maksym (@maksym-mishchenko))) Thanks for doing this very useful feature! Overall looks good to me. Copilot reviews found some interesting cases that are worth addressing.

Hi Anand Krishnamoorthi (Anand Krishnamoorthi (@anakrish)), thanks for the thorough review. I addressed comments, I kept the two API suggestions unchanged for the reasons explained in their threads. Could you please take another look when you have time?

The rationale for both sounds good. Review rerun found 3 more comments worth addressing. Then it should be good to go.

Maksym (@maksym-mishchenko) I see the following drawbacks in this implementation:

• The budget is armed at execute*() time, not at set_memory_budget_config() time (state.rs:19-22 → machine.rs:442-449). set_memory_budget_config sets baseline = 0; active = false; the real baseline is captured on the next execution. So it is a per-execution budget, not a budget window a host can open and close.
• RegoVM::set_data / set_input / set_context / load_program all run before the baseline and are uncharged (stated in docs/limits/memory_budget.md). Via FFI, a host can set a budget, load a 500 MB data blob successfully, and only then have evaluation-on-top bounded. If the goal is per-request isolation, the hostile-input case is the one not covered.

Why haven't you pinned thread level baseline and used every allocation (delta gated) for comparison (just like a global limit)? Why have you chosen execute() calls?

Consider whether the intended primitive is a scope (begin(limit) / end() pinning a thread baseline, consulted by every allocating entry point — data load, compile, execute, serialize) rather than a per-execution budget. set_memory_budget_config is about to become public API with per-execution semantics baked in.

Anand Krishnamoorthi (@anakrish) any thoughts on the above ^?

Apart from the design:

Follow-ups on the fixes from the previous review round

• Amortization was removed rather than refined (machine.rs:673, dispatch.rs:28). Comment 3816390968 argued stride-16 sampling was too coarse and asked for a byte-delta trigger; the fix removed the stride entirely. check_memory_budget() now runs on every instruction, doing a THREAD_COUNTERS.with(...) TLS access per instruction whenever a budget is configured — while the global limit it takes precedence over is still amortized (MEMORY_CHECK_STRIDE = 16 + 32 KiB delta gate). The bench harness gained a regular_memory_budget config; could the numbers be posted?
• The .min() ratchet can make the budget stricter than configured (machine.rs:498). memory_budget_baseline = memory_budget_baseline.min(current) never recovers, so any free of pre-baseline memory observed on the execution thread permanently lowers the baseline. Fail-safe direction, but it produces non-deterministic MemoryBudgetExceeded below the configured limit. release_previous_execution_state() before baseline capture mitigates the main in-engine case. The doc says only "does not grant credit" — it should also state the budget can become stricter than configured.

New

• memory_budget_active is never cleared (machine.rs:448; only reset at machine.rs:431 and on the next execution). After a successful execute() the window stays live, so a Rust host calling the public check_memory_budget() later charges unrelated work to the last execution. Given it is pub and the doc invites bindings to call it post-serialization, either consume the window on read or document it as valid only immediately after execute, on the same thread.
• Binding parity — only C# is wired up. C++ (bindings/cpp/regorus.hpp), Go, Java, Python, Ruby, and WASM get the new RegorusStatus variants but no budget API. Fine as an increment, but worth an explicit note in the PR or a tracking issue.
• Doc overstatement — docs/limits/memory_budget.md says the check covers "native string marshaling" for C#. The managed-string copy happens after the FFI call returns and is on the managed heap, not the mimalloc heap, so it is not covered.

• RegoVM::set_data / set_input / set_context / load_program all run before the baseline and are uncharged (stated in docs/limits/memory_budget.md). Via FFI, a host can set a budget, load a 500 MB data blob successfully, and only then have evaluation-on-top bounded. If the goal is per-request isolation, the hostile-input case is the one not covered.

Mark Birger (@kusha)
I can see rationale for both sides. If the data/input are not included in the evaluation budget then, as you showed, really large objects can escape the memory check.
On the other hand, it is easier for the caller to reason about/estimate the size of input and data (say from the length of the byte streams from which they are originally deserialized) and perform checks whereas it is harder for the caller to reason about a policy evaluation's memory consumption. This PR addresses the latter.

However, I do see that it would be nice to also limit input, data memory consumption.

Some challenges around implementation:

• Setting input, data and evaluation is not an atomic operation. One thread could set input/data and another could perform the evaluation. Per thread budgets don't play nicely with this.
• If input/data exist as C# objects as opposed to a byte stream, then they must anyways be serialized on the C# size..ie the 500MB limit would be exceeded even before reaching rust.

Some ideas

• Should we have individual limits around data/object sizes?
• Or for evaluation, should the base line start at earlier or set data/input and end at the end of evaluation? Ideally in case of FFI after the result has been serialized.

Why haven't you pinned thread level baseline and used every allocation (delta gated) for comparison (just like a global limit)? Why have you chosen execute() calls?

Maybe this would be the simplest approach. Have a thread level limit on alive memory?
But if one thread is the producer of programs and input/data and other threads are merely users, then it may not bounded correctly.

So it is a per-execution budget, not a budget window a host can open and close.

Yes. That was also my impression from the PR description and the implementation. An execution-level budget is the clearest semantic: it bounds the additional working memory used during one RVM execution and intentionally excludes objects such as input, data, and the program that may have been created before—and may outlive—that execution.

However, I do agree that large data/input may go undetected by this budget, even if they are eventually subject to configured global memory limit.

We should meet to discuss your scenario and see if the design can be generalized and the semantics made more clear/easy to reason about.

Mark Birger (@kusha) Anand Krishnamoorthi (@anakrish). I made another pass based on this discussion.

The existing execute* APIs still have a fresh budget for each execution. I added an opt-in path that starts before evaluation-specific JSON data is parsed and set, then stays active through execution and native result serialization.

I avoided a begin/end API across FFI calls because the accounting is thread-local, and separate calls may run on different OS threads. Rust uses a same-thread scope; FFI and C# do everything in one call.

Program loading and compilation are still excluded. I also left input and context out for now. Is covering data enough for the fetch scenario, or do you think input should be included before we approve the API?

The cleanup and documentation issues are fixed as well. The benchmark showed no measurable difference: 859.82 ns without the budget and 847.59 ns with it, with overlapping intervals.

Maksym (@maksym-mishchenko) data is typically the data that doesn't share that often across individual policy evaluation. input is the actual input for a policy evaluation. IMO input would also need to be guarded.

I found a performance issue in the new memory-budget enforcement path: memory_check() is invoked twice per VM instruction in run-to-completion execution. src/rvm/vm/execution.rs calls self.memory_check() before dispatch, and src/rvm/vm/dispatch.rs calls it again at the start of execute_instruction(). The same duplicate exists on the suspendable dispatch path.

With allocator limits enabled, each call can sample the budget and run the throttled global-limit check, so this adds avoidable hot-path overhead and makes checkpoint cadence effectively 2x what the execution loop suggests.

Suggested fix: retain the check at one dispatch boundary (preferably the outer execution loop, unless a specific instruction-level invariant requires the inner check), then add a test-only check counter or benchmark assertion to verify one check per dispatched instruction. Please also compare the benchmark with one vs. two checks so the overhead is measurable.

Documentation suggestion — cooperative enforcement and overshoot:

Please make the hard-cap distinction especially explicit in docs/limits/memory_budget.md and the C# README/API docs. The budget is checked at VM checkpoints and after native result serialization; a single instruction, builtin, to_json_str(), or CString allocation can temporarily exceed the configured limit before the next check. This is a cooperative observed-live-bytes budget, not an allocation-time peak-memory cap. A short example of this behavior would help callers choose an appropriate headroom.

Documentation suggestion — VM reuse and warm pools:

Please document that budget outcomes can depend on prior executions when the same RegoVM is reused. Execution state is reset and the baseline is fresh, but retained VM pools/capacities are already live outside the new baseline, so an identical policy/input may have different allocation behavior on a warm VM than on a fresh VM. This is an important qualification to the phrase “fresh per-execution budget.”

Documentation suggestion — baseline ratcheting and the usage diagnostic:

Please clarify that the baseline is lowered when sampled same-thread live bytes fall below the prior baseline, and that this lost headroom is not restored. As a result, the effective budget can become stricter after unrelated or legitimate same-thread frees, and VmError::MemoryBudgetExceeded { usage, .. } should not be interpreted as exact memory owned by the policy evaluation or as a net-allocation measurement.

Documentation suggestion — host callbacks and cross-thread frees:

Please call out that accounting is based on the current thread’s live-byte counter, not allocation ownership. Allocations/frees performed by synchronous custom or host builtins on the execution thread are therefore observed as part of the budget; objects allocated on one thread and freed on another can temporarily skew per-thread/global observations until counters are published. This would help embedding applications avoid assuming strict query-owned attribution or thread-independent accounting.